# -*- coding: utf-8 -*-
# @Time    : 2023/10/16 17:18
# @Author  : 王摇摆
# @FileName: file1.py
# @Software: PyCharm
# @Blog    : https://blog.csdn.net/weixin_44943389?type=blog

# 导入必要的库
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# 读取数据
data = pd.read_csv('../dataset/train.csv')
test_data = pd.read_csv('../dataset/test.csv')
print('[1. 数据集加载完毕]')

# 分离特征和目标变量
X = data.drop(columns=['id', 'target'])
y = data['target']

test_X = test_data.drop(columns='id')

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

print('[2. 数据集预处理完成]')

# 初始化并训练随机森林模型
random_forest = RandomForestClassifier(n_estimators=100, random_state=42)
random_forest.fit(X_train, y_train)
print('[3. 随机森林模型训练完成]')

# 预测
y_pred = random_forest.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
print(f'随机森林准确度: {accuracy}\n')

# 预测结果输出
test_pred = random_forest.predict(test_X)
pd.DataFrame({'id': test_data['id'], 'target': test_pred}).to_csv('../result/RandomForest1.csv', index=None)
print('[4. 预测结果已输出为CSV文件]')
